#pragma once

#include "NvInferPlugin.h"
#include "cuda_runtime_api.h"

// #include "../byte_tracker/object_detect_info.h"
#include "../pose_detector/pose_util.h"
#include "../trt_infer.h"

#include <rclcpp/rclcpp.hpp>
#include <mutex>

namespace perception::camera {

struct ReIDBinding {
    size_t size = 1;
    size_t dsize = 1;
    nvinfer1::Dims dims;
    std::string name;
};

class ReID {
public:
    bool init(const std::string& engine_file_path);
    void deinit();
    // 返回图特征，已归一化，可以直接送入cosine_dist
    int process(cv::Mat image, std::vector<float>& feat);
    void make_pipe(bool warmup);
    void letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size);
    void copy_from_mat(const cv::Mat& image);
    void infer();
    void postprocess(std::vector<float>& feat);
    // 归一化的向量间，cosine 距离是两个向量的夹角，范围 -1，1
    float cosine_dist(const std::vector<float>& normed_feat1, const std::vector<float>& normed_feat2);

    static std::shared_ptr<ReID> get_detector(const std::string& engine_file_path);

    int num_bindings;
    int num_inputs = 0;
    int num_outputs = 0;
    long long time_total_ = 0;
    long long infer_count_ = 0;
    std::vector<ReIDBinding> input_bindings;
    std::vector<ReIDBinding> output_bindings;
    std::vector<void*> host_ptrs;
    std::vector<void*> device_ptrs;
    std::string cam_name_;

private:
    ReID();
    std::mutex pipe_mutex_;
    cv::Mat image_resize_, image_orig_;
    // 推理 h 实际为384, engine内部做pad, postprocess 需要减去pad
    // 网络要求size 需要为24的整数倍
    int infer_h_ = 128; 
    int infer_w_ = 128;
    int pad_h_ = 12;
    int pad_w_ = 0;
    int iframe_ = 0;
    std::map<int, int> hw2air_;
    nvinfer1::ICudaEngine* engine_ = nullptr;
    nvinfer1::IRuntime* runtime_ = nullptr;
    nvinfer1::IExecutionContext* context_ = nullptr;
    cudaStream_t stream_ = nullptr;
    NvLogger logger_ { nvinfer1::ILogger::Severity::kERROR };
    float* data_;
    static std::shared_ptr<ReID> detector_;

    rclcpp::Logger logger__;
};

}